# 加载飞桨和相关数据处理的库
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Linear
import numpy as np
import os
import gzip  # 打包文件处理库
import json  # JSON 数据格式库
import random  # 随机文件处理库


#? 封装成函数方便功能的使用
def load_data(mode='train'):
    #. 1. 读取文件名称为`mnist.json.gz`的MNIST数据
    #> 声明 datafile 保存数据集文件路径
    datafile = './work/mnist.json.gz'
    # print('mnist 数据路径是：  {} ......'.format(datafile))
    #> 加载json数据文件 （两个库均在上面import项目中）
    #@ 1. gzip.open 打开压缩数据（参数为路径 ）
    #@ 2. json.load 将解压的数据转换成JSON格式数据
    data = json.load(gzip.open(datafile))
    # print('mnist 数据加载完成')
    # print(type(data))
    #! print('预览数据 data', data) 还是不要预览了，程序司机了，内部样子大概是这样的

    #> 读取到的数据区分训练集，验证集，测试集
    #@ 这个语法是一种解包的概念
    """
    data = [ 
          [ 
              [784个数据],[784个数据]....5万个数组], [0,1,2,3,4...5万] 
          ] , 
          [ 
              [784个数据],[784个数据]....1万个数组], [0,1,2,3,4...1万] 
          ] ,
          [ 
              [784个数据],[784个数据]....1万个数组], [0,1,2,3,4...1万] 
          ] ,
      ]
    """
     # 读取到的数据区分训练集，验证集，测试集
    train_set, val_set, eval_set = data
    if mode=='train':
        # 获得训练数据集
        imgs, labels = train_set[0], train_set[1]
    elif mode=='valid':
        # 获得验证数据集
        imgs, labels = val_set[0], val_set[1]
    elif mode=='eval':
        # 获得测试数据集
        imgs, labels = eval_set[0], eval_set[1]
    else:
        raise Exception("mode can only be one of ['train', 'valid', 'eval']")
    print("训练数据集数量: ", len(imgs))

    #> 数据集相关参数，图片高度IMG_ROWS, 图片宽度IMG_COLS
    #@ 常量用全大写字母书写，这是编程习惯
    IMG_ROWS = 28
    IMG_COLS = 28

    #> 获得数据集长度
    imgs_length = len(imgs)
    
    #! 效验数据有效性 1 对比数据和标签数据长度是否一致
    assert len(imgs) == len(
        labels), "train_imgs 的长度 ({}) 应该与train_labels({}) 相等".format(
            len(imgs), len(label))
    
    #> 定义数据集每个数据的序号，根据序号读取数据
    #@ range(5万)生成 0-5万的整数列表 返回的数据类型  <class 'range'>
    #@ list() 强转一次数据类型为列表
    index_list = list(range(imgs_length))
    
    #> 读入数据时用到的批次大小
    BATCHSIZE = 100

    #> 定义数据生成器，返回批次数据
    def data_generator():
        if mode == 'train':
            #> 随机打乱训练数据的索引序号!!!!!!!
            random.shuffle(index_list)
        #> 数据列表
        imgs_list = []
        #> 标签列表
        labels_list = []
        #> 循环索引列表
        for i in index_list:

            #> 将数据处理成期望的格式，比如类型为float32，shape为[1, 28, 28]
            #@ np.reshape 不改变数据内容的情况下，改变一个数组的格式
            #@ 第一个参数 ： 需要处理的数据 =imgs[i]
            #@ 第二个参数： 新的格式 [1, 28, 28]
            #@ 'float32‘ 是数据的小数点有效长度的设置。
            #@ 最终img的数据格式[ [28个数据],[28个数据],...26个列表]
            img = np.reshape(imgs[i],
                             [1, IMG_ROWS, IMG_COLS]).astype('float32')

            #@ labels[0] 是5
            label = np.reshape(labels[i], [1]).astype('float32')
            
            #> 添加img数据到imgs列表
            imgs_list.append(img)
            
            #> 添加标签数据到labels列表
            labels_list.append(label)
            
            #> 如果图片列表的长度 == 100
            if len(imgs_list) == BATCHSIZE:
                #> 获得一个batchsize的数据，并返回 目的是以100个数据为一个批次的返回依靠关键字yield实现
                """
                yield 关键字，为了控制超大数据量循环时造成的内存占用不合理的问题出现的
                详细解释： https://developer.ibm.com/zh/articles/os-cn-python-yield/
                这个最新 ：https://blog.csdn.net/u011318077/article/details/93749143
                这行代码目的可以理解为： 
                    每次data_generator函数被调用一次，则会返回2个数据：
                    1个数据类型为np.array的imgs_list列表数据
                    1个数据类型为np.array的labels_list列表数据
                返回上面的数据后，则停止代码停止执行 ，也就是下面，清空数据读取列表的这行代码不会执行，
                需要等待下一次调用从下面一行代码开始在执行并继续for循环
                """ 
                yield np.array(imgs_list), np.array(labels_list)
                #> 清空数据读取列表--上面说的是我行代码
                imgs_list = []
                labels_list = []

        #> 如果剩余数据的数目小于BATCHSIZE，因为循环，img列表如果没有到达100，则不会清空说明剩余的内容小于100又不会是0
        #> 则剩余数据一起构成一个大小为len(imgs_list)的mini-batch
        if len(imgs_list) > 0:
            yield np.array(imgs_list), np.array(labels_list)
        return data_generator


#> 数据处理部分之后的代码，数据读取的部分调用load_data函数
#> 定义网络结构，同上一节所使用的网络结构
class MNIST(fluid.dygraph.Layer):
    def __init__(self):
        super(MNIST, self).__init__()
        self.fc = Linear(input_dim=784, output_dim=1, act=None)

    def forward(self, inputs):
        inputs = fluid.layers.reshape(inputs, (-1, 784))
        outputs = self.fc(inputs)
        return outputs

#! 异步训练配置，并启动训练过程
with fluid.dygraph.guard():
    model = MNIST()
    model.train()
    #> 调用加载数据的函数
    train_loader = load_data('train')
    #. 创建异步数据读取器
    place = fluid.CPUPlace()
    data_loader = fluid.io.DataLoader.from_generator(capacity=5, return_list=True)
    data_loader.set_batch_generator(train_loader, places=place)
    
    optimizer = fluid.optimizer.SGDOptimizer(learning_rate=0.001, parameter_list=model.parameters())
    EPOCH_NUM = 3
    for epoch_id in range(EPOCH_NUM):
        for batch_id, data in enumerate(data_loader):
            image_data, label_data = data
            image = fluid.dygraph.to_variable(image_data)
            label = fluid.dygraph.to_variable(label_data)
            
            predict = model(image)
            
            loss = fluid.layers.square_error_cost(predict, label)
            avg_loss = fluid.layers.mean(loss)
            
            if batch_id % 200 == 0:
                print("epoch: {}, batch: {}, loss is: {}".format(epoch_id, batch_id, avg_loss.numpy()))
            
            avg_loss.backward()
            optimizer.minimize(avg_loss)
            model.clear_gradients()

    fluid.save_dygraph(model.state_dict(), 'mnist3_2')
